{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "from torchvision import datasets, transforms\n",
    "import torch.multiprocessing as mp\n",
    "import argparse\n",
    "\n",
    "class MNISTNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MNISTNET,self).__init__()\n",
    "        self.layer1 = nn.Linear(784,50)\n",
    "        self.layer2 = nn.Linear(50,10)\n",
    "    def forward(self,x):\n",
    "        #输入层到隐藏层，使用tanh激活函数\n",
    "        x = self.layer1(x)\n",
    "        x = F.tanh(x)\n",
    "        #隐藏层到输出层，使用relu激活函数\n",
    "        x = self.layer2(x)\n",
    "        x = F.relu(x)\n",
    "        #log(softmax)操作，使用NLLLoss损失函数\n",
    "        x = nn.log_softmax(x)\n",
    "\n",
    "def train(rank, args, model, device, dataloader_kwargs):\n",
    "    #手动设置随机种子\n",
    "    torch.manual_seed(args.seed + rank)\n",
    "    #加载训练数据\n",
    "    train_loader = torch.utils.data.DataLoader(datasets.MNIST('../data', train=True, download=True,transform=transforms.Compose([\n",
    "                        transforms.ToTensor(),\n",
    "                        transforms.Normalize((0.,), (1.,))\n",
    "                    ])),batch_size=args.batch_size, shuffle=True, num_workers=1,**dataloader_kwargs)\n",
    "    #使用随机梯度下降进行优化\n",
    "    optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)\n",
    "    #开始训练，训练epoches次\n",
    "    for epoch in range(1, args.epochs + 1):\n",
    "        train_epoch(epoch, args, model, device, train_loader, optimizer)\n",
    "\n",
    "def test(args, model, device, dataloader_kwargs):\n",
    "    #设置随机种子\n",
    "    torch.manual_seed(args.seed)\n",
    "    #加载测试数据\n",
    "    test_loader = torch.utils.data.DataLoader(datasets.MNIST('../data', train=False, transform=transforms.Compose([\n",
    "            transforms.ToTensor(),\n",
    "            transforms.Normalize((0.,), (1.,))\n",
    "        ])),batch_size=args.batch_size, shuffle=True, num_workers=1,**dataloader_kwargs)\n",
    "    #运行测试\n",
    "    test_epoch(model, device, test_loader)\n",
    "    \n",
    "    \n",
    "def train_epoch(epoch, args, model, device, data_loader, optimizer):\n",
    "    #模型转换为训练模式\n",
    "    model.train()\n",
    "    pid = os.getpid()\n",
    "    for batch_idx, (data, target) in enumerate(data_loader):\n",
    "        #优化器梯度置0\n",
    "        optimizer.zero_grad()\n",
    "        #输入特征预测值\n",
    "        output = model(data.to(device))\n",
    "        #预测值与标准值计算损失\n",
    "        loss = F.nll_loss(output, target.to(device))\n",
    "        #计算梯度\n",
    "        loss.backward()\n",
    "        #更新梯度\n",
    "        optimizer.step()\n",
    "        #每10步打印一下日志\n",
    "        if batch_idx % 10 == 0:\n",
    "            print('{}\\tTrain Epoch: {} [{}/{} ({:.0f}%)]\\tLoss: {:.6f}'.format(pid, epoch, batch_idx * len(data), len(data_loader.dataset),\n",
    "                100. * batch_idx / len(data_loader), loss.item()))\n",
    "\n",
    "\n",
    "def test_epoch(model, device, data_loader):\n",
    "    #模型转换为测试模式\n",
    "    model.eval()\n",
    "    test_loss = 0\n",
    "    correct = 0\n",
    "    with torch.no_grad():\n",
    "        for data, target in data_loader:\n",
    "            output = model(data.to(device))\n",
    "            #将每个批次的损失加起来\n",
    "            test_loss += F.nll_loss(output, target.to(device), reduction='sum').item()\n",
    "            #得到概率最大的索引,\n",
    "            pred = output.max(1)[1]\n",
    "            #预测的索引和目标索引相同，认为预测正确\n",
    "            correct += pred.eq(target.to(device)).sum().item()\n",
    "\n",
    "    test_loss /= len(data_loader.dataset)\n",
    "    print('\\nTest set: Average loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\\n'.format(est_loss, correct, len(data_loader.dataset),\n",
    "        100. * correct / len(data_loader.dataset)))\n",
    "\n",
    "# 解析传入的参数\n",
    "parser = argparse.ArgumentParser(description='PyTorch MNIST')\n",
    "parser.add_argument('--batch-size', type=int, default=64, metavar='N',\n",
    "                    help='input batch size for training (default: 64)')\n",
    "parser.add_argument('--test-batch-size', type=int, default=1000, metavar='N',\n",
    "                    help='input batch size for testing (default: 1000)')\n",
    "parser.add_argument('--epochs', type=int, default=10, metavar='N',\n",
    "                    help='number of epochs to train (default: 10)')\n",
    "parser.add_argument('--lr', type=float, default=0.01, metavar='LR',\n",
    "                    help='learning rate (default: 0.01)')\n",
    "parser.add_argument('--momentum', type=float, default=0.9, metavar='M',\n",
    "                    help='SGD momentum (default: 0.9)')\n",
    "parser.add_argument('--seed', type=int, default=1, metavar='S',\n",
    "                    help='random seed (default: 1)')\n",
    "parser.add_argument('--log-interval', type=int, default=10, metavar='N',\n",
    "                    help='how many batches to wait before logging training status')\n",
    "parser.add_argument('--num-processes', type=int, default=10, metavar='N',\n",
    "                    help='how many training processes to use (default: 10)')\n",
    "parser.add_argument('--cuda', action='store_true', default=False,\n",
    "                    help='enables CUDA training')\n",
    "if __name__==\"__main__\":\n",
    "    #解析参数\n",
    "    args = parser.parse_args()\n",
    "    #判断是否使用GPU\n",
    "    use_cuda = args.cuda and torch.cuda.is_available()\n",
    "    #运行时设备\n",
    "    device = torch.device(\"cuda\" if use_cuda else \"cpu\")\n",
    "    #使用固定缓冲区\n",
    "    dataloader_kwargs = {'pin_memory': True} if use_cuda else {}\n",
    "    #多进程训练，windows使用spawn方式\n",
    "    mp.set_start_method('spawn')\n",
    "    #模型拷贝到设备\n",
    "    model = MNISTNet().to(device)\n",
    "    #多进程共享模型参数\n",
    "    model.share_memory()\n",
    "\n",
    "    processes = []\n",
    "    for rank in range(args.num_processes):\n",
    "        p = mp.Process(target=train, args=(rank,args, model,device, dataloader_kwargs))\n",
    "        p.start()\n",
    "        processes.append(p)\n",
    "    for p in processes:\n",
    "        p.join()\n",
    "\n",
    "    #测试模型\n",
    "    test(args, model, device, dataloader_kwargs)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([32, 1, 28, 28])"
      ]
     },
     "execution_count": 38,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "data[0].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
